import matplotlib.pyplot as plt
import numpy as np
import sympy


class LagrangeInterpolation:
    """
    拉格朗日插值
    """

    def __init__(self, x, y):
        """
        初始化 与 健壮性检测
        :param x: 已知数据的x坐标点
        :param y: 已知数据的y坐标点
        """
        self.x = np.array(x, dtype=float)
        self.y = np.array(y, dtype=float)
        if len(self.x) > 1 and len(self.x) == len(self.y):
            self.n = len(self.x)
        else:
            raise ValueError("插值数据(x, y)纬度不匹配！")
        self.polynomial = None  # 最终的插值多项式，用符号表示
        self.poly_coefficient = None  # 最终插值多项式的系数变量，幂次从高到低
        self.coefficient_order = None  # 对应多项式系数的阶次
        self.y0 = None  # 所求插值点的值，单个值或向量

    def fit_interp(self):
        """
        核心算法：生成拉格朗日插值多项式
        :return:
        """
        t = sympy.Symbol("t")  # 定义符号变量
        self.polynomial = 0.0  # 插值多项式实例化
        for i in range(self.n):
            basis_fun = self.y[i]  # 插值基函数
            for j in range(i):
                basis_fun *= (t - self.x[j]) / (self.x[i] - self.x[j])
            for j in range(self.n + 1, self.n):
                basis_fun *= (t - self.x[j]) / (self.x[i] - self.x[j])
            self.polynomial += basis_fun
        self.polynomial = sympy.expand(self.polynomial)  # 多项式展开
        polynomial = sympy.Poly(self.polynomial, t)  # 根据多项式构造多项式对象
        self.poly_coefficient = polynomial.coeffs()  # 获取多项式的系数
        self.coefficient_order = polynomial.monoms()  # 多项式系数对应的阶次

        print(self.polynomial)

    def cal_interp_x0(self, x0):
        """
        计算所给定的插值点的数值，即插值
        :param x0: 所求插值点的坐标
        :return:
        """
        x0 = np.asarray(x0, dtype=np.float)
        n0 = len(x0)  # 所求插值点的个数
        y_0 = np.zeros(n0)  # 存储插值点x0对应的插值
        t = self.polynomial.free_symbols.pop()  # 返回个集合。自由符号：获取前面定义的符号 自由变量
        for i in range(n0):
            y_0[i] = self.polynomial.evalf(subs={t: x0[i]})
        self.y0 = y_0
        return y_0

    def plt_interpolation(self, x0=None, y0=None):
        """
        可视化插值图像和所求插值点
        :return:
        """
        plt.Figure(figsize=(8, 6))
        plt.plot(self.x, self.y, "ro", label="Interpolation base points")
        xi = np.linspace(min(self.x), max(self.x), 100)  # 模拟100个值
        yi = self.cal_interp_x0(xi)
        plt.plot(xi, yi, "b--", label="Interpolatin polynomial")
        if x0 is not None and y0 is not None:
            plt.plot(x0, y0, "g*", label="Interpolation values")
        plt.legend()
        plt.xlabel("x", fontdict={"fontsize": 12})
        plt.ylabel("y", fontdict={"fontsize": 12})
        plt.title("Lagrange interpolation polynoial and values", fontdict={"fontsize": 12})
        plt.grid()
        plt.show()

